import pandas as pd
import torch
import torch.nn as nn
from d2l import torch as d2l

from utils.file_download_util import FileDownloader


# 训练集的对数均方根误差（Root Mean Square Error，RMSE）
# 这个函数接收三个参数：net 是一个神经网络模型，features 是输入特征数据，labels 是对应的真实标签数据
def log_rmse(net, features, labels):
    # 对于一个net的输出如果是inf的话就将其变为1
    clipped_preds = torch.clamp(net(features), 1, float('inf'))
    rmse = torch.sqrt(loss(torch.log(clipped_preds), torch.log(labels)))
    return rmse.item()  # 取出单元素张量的元素值并返回该值，保持原元素类型不变


# 构建一个单层神经网络
def get_default_net():
    net = nn.Sequential(nn.Linear(in_features, 1))
    return net


# 多个模型的工具调用方法
def train(net, train_features, train_labels, test_features, test_labels,
          num_epochs, learning_rate, weight_decay, batch_size):
    """
    训练神经网络的函数。

    参数:
    net -- 要训练的神经网络模型
    train_features -- 训练数据的特征
    train_labels -- 训练数据的标签
    test_features -- 测试数据的特征（可选）
    test_labels -- 测试数据的标签（可选）
    num_epochs -- 训练的轮数
    learning_rate -- 优化器的学习率
    weight_decay -- 权重衰减，用于正则化
    batch_size -- 每个批次的数据量

    返回:
    train_ls -- 训练损失的列表
    test_ls -- 测试损失的列表（如果提供了测试数据）
    """
    # 如果未传递 net 参数，则使用 get_default_net() 获取默认网络模型
    if net is None:
        net = get_default_net()
    train_ls, test_ls = [], []  # 存储每个epoch结束时的训练和测试损失
    train_iter = d2l.load_array((train_features, train_labels), batch_size)  # 创建数据迭代器
    optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=weight_decay)  # 创建优化器

    for epoch in range(num_epochs):  # 遍历每个epoch
        for X, y in train_iter:  # 遍历训练数据的每个批次
            optimizer.zero_grad()  # 清零梯度
            l = loss(net(X), y)  # 计算损失（这里没有使用对数均方误差，可能是因为训练时关注的是原始损失的梯度）
            l.backward()  # 反向传播
            optimizer.step()  # 更新权重
        # 在每个epoch结束后记录训练集的对数均方根误差（log_rmse）
        train_ls.append(log_rmse(net, train_features, train_labels))
        if test_labels is not None:  # 如果提供了测试数据
            # 计算并记录测试集的对数均方根误差
            test_ls.append(log_rmse(net, test_features, test_labels))
    return train_ls, test_ls  # 返回训练和测试损失的列表


# 调用k_fold函数进行训练和验证，获取每一折的具体训练和验证数据
def get_k_fold_data(k, i, X, y):
    assert k > 1  # k肯定大于1
    fold_size = X.shape[0] // k
    X_train, y_train = None, None
    for j in range(k):
        # 切片，nep[slice(2,4)] == nep[2:4]
        idx = slice(j * fold_size, (j + 1) * fold_size)
        # 取出当前折
        X_part, y_part = X[idx, :], y[idx]
        if j == i:  # vaildation set
            X_vaild, y_vaild = X_part, y_part
        elif X_train is None:
            # 为空的话就先赋值
            X_train, y_train = X_part, y_part
        else:
            # 非空的话就concat
            X_train = torch.cat([X_train, X_part], 0)
            y_train = torch.cat([y_train, y_part], 0)
    return X_train, y_train, X_vaild, y_vaild


# 调用k_fold函数进行训练和验证，返回训练损失和验证损失
def k_fold(k, X_train, y_train, num_epochs, learning_rate, weight_decay, batch_size):
    train_ls_sum, vaild_ls_sum = 0, 0
    # 进行k折交叉验证
    for i in range(k):
        data = get_k_fold_data(k, i, X_train, y_train)
        net = get_default_net()
        # *data 表示打包非键值对参数data，对应tupple (训练特征，训练标签， 测试特征， 测试标签)
        # trian()返回的是tupple (训练误差， 测试误差) 可以直接像这样用两个元素去接
        train_ls, valid_ls = train(net, *data, num_epochs, learning_rate, weight_decay, batch_size)
        train_ls_sum += train_ls[-1]  # -1表示取这一轮最后一个epoch的loss
        vaild_ls_sum += valid_ls[-1]
        if i == 0:
            # 自变量 因变量 标签  x轴 图例说明  y轴
            d2l.plot(list(range(1, num_epochs + 1)), [train_ls, valid_ls],
                     xlabel='epoch', ylabel='rmse', xlim=[1, num_epochs],
                     legend=['train', 'valid'], yscale='log')
        print(f'Nep！fold{i + 1},train log rmse {float(train_ls[-1]):f}, '
              f'valid log rmse {float(valid_ls[-1]):f}')
    return train_ls_sum / k, vaild_ls_sum / k


# 按间距中的绿色按钮以运行脚本。
if __name__ == '__main__':
    print('----------')

    # 1-定义要下载的文件
    DATA_URL = 'http://d2l-data.s3-accelerate.amazonaws.com/'
    FileDownloader.DATA_HUB = {
        # 二元组分别对应其url和密钥
        'kaggle_house_train': (DATA_URL + 'kaggle_house_pred_train.csv', '585e9cc93e70b39160e7921475f9bcd7d31219ce'),
        'kaggle_house_test': (DATA_URL + 'kaggle_house_pred_test.csv', 'fa19780a7b011d9b009e8bff8e99922a8ee2eb90'),
    }

    # 2-进行文件下载
    train_data = pd.read_csv(FileDownloader.download('kaggle_house_train'))
    test_data = pd.read_csv(FileDownloader.download('kaggle_house_test'))

    # 3-查看数据格式
    print('----------train_data.shape：', train_data.shape)
    print('----------test_data.shape：', test_data.shape)

    # 4-查看前四个样本的前四个后最后两个特征，以及相应的标价
    print('----------train_data.iloc：', train_data.iloc[0:4, [0, 1, 2, 3, -3, -2, -1]])

    # 5-【数据预处理】我们将其id信息从数据集中删除
    # 去掉ID后的所有样本,都表示对于训练集/测试集的所有样本，从下标1一直到最后，使用pd的concat将二者合并
    all_features = pd.concat((train_data.iloc[:, 1:-1], test_data.iloc[:, 1:]))
    print('----------all_features.shape：', all_features.shape)

    # 6-【数据预处理-数值特征】将所有缺失值替换为相应特征的平均值。然后进行feature scaling（先标准化再替换会更方便，因为标准化后均值就是0了）
    # 假设不是'object'类的数据就是数值类，把所有数值特征提取出来
    numeric_feaures = all_features.select_dtypes(include=['number']).columns
    # 对数值特征进行标准化(实际情况拿不到测试集的时候可以用训练集的均值方差代替)
    all_features[numeric_feaures] = all_features[numeric_feaures].apply(
        lambda x: (x - x.mean()) / (x.std()))
    # 把NAN变成均值，不建议改成0
    all_features[numeric_feaures] = all_features[numeric_feaures].fillna(all_features[numeric_feaures].mean())

    # 7-【数据预处理-离散值特征】进行one-hot独热编码替换
    all_features = pd.get_dummies(all_features, dummy_na=True)
    # 将布尔值转换为0和1(当前Python版本默认使用True/False)
    all_features = all_features.astype(int)
    # 填充缺失值为相应列的均值
    all_features = all_features.fillna(all_features.mean())

    # 8-【数据预处理-分割数据集】将数据集分割为训练集和测试集
    # 转换为torch张量
    n_train = train_data.shape[0]
    # 从all_features中取出前n_train行作为训练集，后n_test行作为测试集
    train_features = torch.tensor(all_features[:n_train].values, dtype=torch.float32)
    test_features = torch.tensor(all_features[n_train:].values, dtype=torch.float32)
    # 从train_data中取出前n_train行作为训练集，后n_test行作为测试集
    train_labels = torch.tensor(train_data.SalePrice.values, dtype=torch.float32).view(-1, 1)

    # 9-【数据预处理-查看数据】查看数据
    print('----------train_features.shape：', train_features.shape)
    print('----------train_labels.shape：', train_labels.shape)
    print('----------test_features.shape：', test_features.shape)

    # 10-【模型训练】训练⼀个带有损失平⽅的线性模型
    loss = nn.MSELoss()
    in_features = train_features.shape[-1]  # 提取所有训练特征

    # 从文件: main07_单个模型训练_文件上传.py
    # 设置超参数
    # k折交叉验证的折数，训练轮数，学习率，权重衰减，批量大小
    k, num_epochs, lr, weight_decay, batch_size = 5, 100, 5, 0, 64
    # 调用k_fold函数进行训练和验证，返回训练损失和验证损失
    train_l, valid_l = k_fold(k, train_features, train_labels, num_epochs, lr, weight_decay, batch_size)

    print('k折交叉验证全部搞定！')
    print(f'{k}折验证：平均训练log rmse：{float(train_l):f},'
          f'平均验证log rmse：{float(valid_l):f}')
